#include <stdio.h>
#include <iostream>
#include <fstream>
#include <log.hpp>
#include <memfs.h>

CLog clg;

typedef struct loc_s
{
  int type;
  int processor;
} loc_t;

int getoperations(Mshell &trace, loc_t **operations, int *numops);
int colocation(loc_t *operations, int numops, int **types, int *numtypes, int **affinity, int *numprocs);
int printcolocation(int *types, int numtypes);
int printaffinity(int *affinity, int numtypes, int numprocs);
int outputcolocation(int *types, int numtypes, Mshell &coloc);
int outputaffinity(int *affinity, int numtypes, int numprocs, Mshell &coloc);

extern "C" int MAIN(int argc, char **argv)
{
  Mfs tfs, cfs;
  Mshell trace(tfs), coloc(cfs);
  loc_t *operations;
  int numops, *types, numtypes, *affinity, numprocs;

  if (argc != 2 && argc != 3)
  {
    clg.err() << "Usage: " << argv[0] << " <trace> [colocations]" << std::endl;
    return -1;
  }

  std::ifstream tif(argv[1]);
  XMLParser txp(tif);
  trace.Import(txp, "/");

  getoperations(trace, &operations, &numops);
  colocation(operations, numops, &types, &numtypes, &affinity, &numprocs);

  if (argc == 2)
  {
    printcolocation(types, numtypes);
    printaffinity(affinity, numtypes, numprocs);
  }
  else
  {
    outputcolocation(types, numtypes, coloc);
    outputaffinity(affinity, numtypes, numprocs, coloc);

    std::ofstream cof(argv[2]);
    XMLParser cxp(cof);
    coloc.Export(cxp, "/");
  }

  return 0;
}

int getoperations(Mshell &trace, loc_t **operations, int *numops)
{
  char *element;

  *numops = 0;
  *operations = NULL;

  trace.FindInitDirectoryElements("/trace/operations");
  while (!trace.FindGetNextElement(&element))
  {
    if (!((*numops)%100))
      *operations = (loc_t*) realloc(*operations, ((*numops)+100)*sizeof(loc_t));

    trace.ChDir(element);

    if (trace.GetContent("type", &(*operations)[*numops].type))
      clg.err() << HDR << "Couldn't get operation type for '" << element << "'" << std::endl;

    if (trace.GetContent("processor", &(*operations)[*numops].processor))
      clg.err() << HDR << "Couldn't get processor for '" << element << "'" << std::endl;

    (*numops)++;
  }

  return 0;
}

int colocation(loc_t *operations, int numops, int **types, int *numtypes,
                                              int **affinity, int *numprocs)
{
  int *used;

  *numtypes = 0;
  *numprocs = 0;

  for (int i=0; i < numops; i++)
  {
    if (operations[i].type > *numtypes)
      *numtypes = operations[i].type;
    if (operations[i].processor > *numprocs)
      *numprocs = operations[i].processor;
  }

  (*numtypes)++;
  (*numprocs)++;

  *types = (int*) calloc((*numtypes)*(*numtypes), sizeof(int));
  *affinity = (int*) calloc((*numtypes)*(*numprocs), sizeof(int));
  used = (int*) malloc((*numtypes)*sizeof(int));

  for (int i=0; i < numops; i++)
  {
    for (int j=0; j < (*numtypes); j++)
      used[j] = 0;

    (*affinity)[operations[i].type * (*numprocs) + operations[i].processor]++;

    /* Find operations which follow this one. For every type, the first following
     * operation defines the colocation */
    for (int j=i; j < numops; j++)
      if (!used[operations[j].type])
      {
	used[operations[j].type] = true;
	if (operations[i].processor == operations[j].processor)
	  (*types)[operations[i].type * (*numtypes) + operations[j].type]++;
      }
    
    /* Operations which do not follow have 100% colocation automatically */
/*
    for (int j=0; j < numops; j++)
    {
      if (!used[operations[j].type])
      {
	used[operations[j].type] = true;
	(*types)[operations[i].type * (*numtypes) + operations[j].type]++;
      }
    }
*/
  }

  free(used);

  return 0;
}

int printcolocation(int *types, int numtypes)
{
  printf("   ");
  for (int j=0; j < numtypes; j++)
  {
    if (types[j*numtypes+j])
      printf("%4d ", j);
  }
  printf("\n");

  for (int i=0; i < numtypes; i++)
  {
    if (types[i*numtypes+i])
    {
      printf("%2d ", i);

      for (int j=0; j < numtypes; j++)
      {
        if (i==j)
        {
          printf("  .  ");
        }
        else if (types[j*numtypes+j])
	{
	  if (types[i*numtypes+j] == types[i*numtypes+i])
	    printf("  *  ");
	  else if (types[i*numtypes+j])
	    printf("%4.2f ", types[i*numtypes+j]/(double)types[i*numtypes+i]);
	  else
	    printf("     ");
	}
      }

      printf("\n");
    }
  }

  return 0;
}

int printaffinity(int *affinity, int numtypes, int numprocs)
{
  int *pused, *tused;

  tused = (int*) calloc(numtypes, sizeof(int));
  pused = (int*) calloc(numprocs, sizeof(int));

  for (int i=0; i < numtypes; i++)
    for (int j=0; j < numprocs; j++)
    {
      tused[i] += affinity[i * numprocs + j];
      pused[j] += affinity[i * numprocs + j];
    }

  printf("   ");
  for (int j=0; j < numprocs; j++)
  {
    if (pused[j])
      printf("%4d ", j);
  }
  printf("\n");

  for (int i=0; i < numtypes; i++)
  {
    if (tused[i])
    {
      printf("%2d ", i);

      for (int j=0; j < numprocs; j++)
      {
        if (pused[j])
        { 
	  if (affinity[i * numprocs + j])
	    printf("%4.2f ", affinity[i*numprocs+j]/(double)tused[i]);
	  else
	    printf("     ");
	}
      }

      printf("\n");
    }
  }

  free(tused);
  free(pused);

  return 0;
}

int outputcolocation(int *types, int numtypes, Mshell &coloc)
{
  char buf[PATH_MAX];

  for (int i=0; i < numtypes; i++)
  {
    for (int j=0; j < numtypes; j++)
    {
      sprintf(buf, "document/colocation/op%d/op%d", i, j);

      if (types[i*numtypes+i])
	coloc.MkEntry(buf, types[i*numtypes+j]/(double)types[i*numtypes+i]);
      else
	coloc.MkEntry(buf, 0.0);
    }
  }

  return 0;
}

int outputaffinity(int *affinity, int numtypes, int numprocs, Mshell &coloc)
{
  char buf[PATH_MAX];
  int *used;

  used = (int*) calloc(numtypes, sizeof(int));

  for (int i=0; i < numtypes; i++)
    for (int j=0; j < numprocs; j++)
      used[i] += affinity[i * numprocs + j];

  for (int i=0; i < numtypes; i++)
  {
    for (int j=0; j < numprocs; j++)
    {
      sprintf(buf, "document/affinity/op%d/op%d", i, j);

      if (affinity[i*numprocs+j])
	coloc.MkEntry(buf, affinity[i*numprocs+j]/(double)used[i]);
      else
	coloc.MkEntry(buf, 0.0);
    }
  }

  return 0;
}
